"""
MySQL数据库封装模块
用于保存AI生成的剧场数据到本地数据库
数据库: theam, 用户: root, 密码: root
"""

import json
import mysql.connector
from mysql.connector import Error
from typing import Dict, List, Any, Optional, Tuple
from datetime import datetime
import logging
from contextlib import contextmanager
import asyncio
from concurrent.futures import ThreadPoolExecutor

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class TheaterMySQLManager:
    """剧场数据MySQL管理器"""
    
    def __init__(self, host='localhost', database='theam', user='root', password='root', port=3306):
        self.config = {
            'host': host,
            'database': database,
            'user': user,
            'password': password,
            'port': port,
            'charset': 'utf8mb4',
            'collation': 'utf8mb4_unicode_ci',
            'autocommit': False,
            'raise_on_warnings': False  # 改为False避免警告
        }
        self.executor = ThreadPoolExecutor(max_workers=3)
        
        # 安全初始化
        try:
            self._safe_init_database()
        except Exception as e:
            logger.warning(f"⚠️ 数据库初始化警告: {e}")
            # 不抛出异常，允许程序继续运行

    def _safe_init_database(self):
        """安全的数据库初始化"""
        try:
            # 测试连接
            self._test_connection()
            # 确保数据库存在
            self._ensure_database_exists()
            # 确保表结构存在
            self._ensure_tables_exist()
            logger.info("✅ 数据库初始化完成")
        except Exception as e:
            logger.error(f"❌ 数据库初始化失败: {e}")
            # 不抛出异常，让程序能够继续运行
            pass

    def _test_connection(self):
        """测试MySQL连接"""
        try:
            # 先测试不指定数据库的连接
            temp_config = self.config.copy()
            temp_config.pop('database')
            
            connection = mysql.connector.connect(**temp_config)
            if connection.is_connected():
                connection.close()
                logger.info("✅ MySQL服务器连接测试成功")
            else:
                raise Exception("连接失败")
        except Exception as e:
            logger.error(f"❌ MySQL连接失败: {e}")
            raise

    def _ensure_database_exists(self):
        """确保数据库存在"""
        temp_config = self.config.copy()
        temp_config.pop('database')
        
        try:
            connection = mysql.connector.connect(**temp_config)
            cursor = connection.cursor()
            
            # 检查数据库是否存在
            cursor.execute("SHOW DATABASES LIKE %s", (self.config['database'],))
            exists = cursor.fetchone()
            
            if not exists:
                # 创建数据库
                cursor.execute(f"CREATE DATABASE `{self.config['database']}` CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
                connection.commit()
                logger.info(f"🗄️ 数据库 '{self.config['database']}' 创建成功")
            else:
                logger.info(f"🗄️ 数据库 '{self.config['database']}' 已存在")
                
        except Exception as e:
            logger.error(f"❌ 数据库操作失败: {e}")
            raise
        finally:
            if cursor:
                cursor.close()
            if connection:
                connection.close()

    def _ensure_tables_exist(self):
        """确保所有表存在"""
        
        tables_sql = {
            "theaters": """
                CREATE TABLE IF NOT EXISTS theaters (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    title VARCHAR(255) NOT NULL COMMENT '剧场标题',
                    description TEXT COMMENT '剧场描述',
                    era VARCHAR(100) COMMENT '时代背景',
                    location VARCHAR(200) COMMENT '故事地点',
                    plot_summary TEXT COMMENT '情节摘要',
                    features JSON COMMENT '特色列表',
                    author VARCHAR(100) DEFAULT 'AI智能生成器' COMMENT '作者',
                    task_id VARCHAR(50) UNIQUE COMMENT 'AI生成任务ID',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '创建时间',
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
                    INDEX idx_title (title),
                    INDEX idx_task_id (task_id),
                    INDEX idx_created_at (created_at)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='剧场主表'
            """,
            
            "characters": """
                CREATE TABLE IF NOT EXISTS characters (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    theater_id INT NOT NULL COMMENT '剧场ID',
                    character_id VARCHAR(50) NOT NULL COMMENT '角色标识ID',
                    name VARCHAR(100) NOT NULL COMMENT '角色名称',
                    description TEXT COMMENT '角色描述',
                    avatar VARCHAR(500) COMMENT '头像URL',
                    selected_skills JSON COMMENT '选择的技能列表',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    INDEX idx_theater_id (theater_id),
                    INDEX idx_character_id (character_id),
                    INDEX idx_name (name)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='角色表'
            """,
            
            "script_nodes": """
                CREATE TABLE IF NOT EXISTS script_nodes (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    theater_id INT NOT NULL COMMENT '剧场ID',
                    node_id VARCHAR(50) NOT NULL COMMENT '节点标识ID',
                    title VARCHAR(255) NOT NULL COMMENT '节点标题',
                    content LONGTEXT COMMENT '节点内容(HTML格式)',
                    character_settings JSON COMMENT '角色设置',
                    level INT DEFAULT 0 COMMENT '节点级别',
                    node_order INT DEFAULT 0 COMMENT '节点顺序',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    INDEX idx_theater_id (theater_id),
                    INDEX idx_node_id (node_id),
                    INDEX idx_order (node_order),
                    INDEX idx_level (level)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='剧本节点表'
            """,
            
            "skills_pool": """
                CREATE TABLE IF NOT EXISTS skills_pool (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    theater_id INT NOT NULL COMMENT '剧场ID',
                    skill_name VARCHAR(100) NOT NULL COMMENT '技能名称',
                    skill_order INT DEFAULT 0 COMMENT '技能顺序',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    INDEX idx_theater_id (theater_id),
                    INDEX idx_skill_name (skill_name),
                    UNIQUE KEY unique_theater_skill (theater_id, skill_name)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='技能池表'
            """,
            
            "generation_logs": """
                CREATE TABLE IF NOT EXISTS generation_logs (
                    id INT AUTO_INCREMENT PRIMARY KEY,
                    task_id VARCHAR(50) NOT NULL COMMENT '任务ID',
                    theater_id INT COMMENT '剧场ID',
                    stage VARCHAR(50) COMMENT '生成阶段',
                    progress DECIMAL(5,2) DEFAULT 0 COMMENT '进度百分比',
                    message TEXT COMMENT '日志消息',
                    error_message TEXT COMMENT '错误消息',
                    metadata JSON COMMENT '额外元数据',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    INDEX idx_task_id (task_id),
                    INDEX idx_theater_id (theater_id),
                    INDEX idx_stage (stage),
                    INDEX idx_created_at (created_at)
                ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_unicode_ci COMMENT='生成日志表'
            """
        }

        try:
            with self.get_connection() as connection:
                cursor = connection.cursor()
                
                for table_name, sql in tables_sql.items():
                    try:
                        cursor.execute(sql)
                        logger.info(f"✅ 表 '{table_name}' 检查完成")
                    except mysql.connector.Error as e:
                        if e.errno in [1050, 1061]:  # 表已存在或索引已存在
                            logger.info(f"📋 表 '{table_name}' 已存在")
                        else:
                            logger.warning(f"⚠️ 表 '{table_name}' 处理警告: {e}")
                
                # 尝试添加外键约束（如果不存在）
                self._add_foreign_keys_if_not_exist(cursor)
                
                connection.commit()
                logger.info("🎯 所有表结构检查完成")
                
        except Exception as e:
            logger.error(f"❌ 创建表失败: {e}")
            raise

    def _add_foreign_keys_if_not_exist(self, cursor):
        """添加外键约束（如果不存在）"""
        foreign_keys = [
            {
                "table": "characters",
                "constraint": "fk_characters_theater",
                "sql": "ALTER TABLE characters ADD CONSTRAINT fk_characters_theater FOREIGN KEY (theater_id) REFERENCES theaters(id) ON DELETE CASCADE"
            },
            {
                "table": "script_nodes", 
                "constraint": "fk_script_nodes_theater",
                "sql": "ALTER TABLE script_nodes ADD CONSTRAINT fk_script_nodes_theater FOREIGN KEY (theater_id) REFERENCES theaters(id) ON DELETE CASCADE"
            },
            {
                "table": "skills_pool",
                "constraint": "fk_skills_pool_theater", 
                "sql": "ALTER TABLE skills_pool ADD CONSTRAINT fk_skills_pool_theater FOREIGN KEY (theater_id) REFERENCES theaters(id) ON DELETE CASCADE"
            }
        ]
        
        for fk in foreign_keys:
            try:
                cursor.execute(fk["sql"])
                logger.info(f"✅ 外键约束 '{fk['constraint']}' 添加成功")
            except mysql.connector.Error as e:
                if e.errno == 1061:  # 重复的键名
                    logger.info(f"📋 外键约束 '{fk['constraint']}' 已存在")
                else:
                    logger.warning(f"⚠️ 外键约束 '{fk['constraint']}' 添加失败: {e}")

    @contextmanager
    def get_connection(self):
        """获取数据库连接的上下文管理器"""
        connection = None
        try:
            connection = mysql.connector.connect(**self.config)
            yield connection
        except Error as e:
            if connection:
                connection.rollback()
            logger.error(f"❌ 数据库连接错误: {e}")
            raise
        finally:
            if connection and connection.is_connected():
                connection.close()

    def save_theater_data(self, task_id: str, theater_data: Dict[str, Any]) -> Dict[str, Any]:
        """保存完整的剧场数据"""
        start_time = datetime.now()
        
        try:
            with self.get_connection() as connection:
                cursor = connection.cursor()
                
                try:
                    # 1. 保存剧场主信息
                    theater_info = theater_data.get('theater_info', {})
                    theater_id = self._save_theater_info(cursor, task_id, theater_info)
                    
                    # 2. 保存角色数据
                    characters_data = theater_data.get('characters', [])
                    character_count = self._save_characters(cursor, theater_id, characters_data)
                    
                    # 3. 保存剧本节点
                    script_nodes = theater_data.get('script_nodes', [])
                    nodes_count = self._save_script_nodes(cursor, theater_id, script_nodes)
                    
                    # 4. 保存技能池
                    available_skills = theater_data.get('available_skills', [])
                    skills_count = self._save_skills_pool(cursor, theater_id, available_skills)
                    
                    # 5. 记录成功日志
                    self._log_generation_result(cursor, task_id, theater_id, "SUCCESS", 
                                              "数据保存成功", {
                                                  "character_count": character_count,
                                                  "nodes_count": nodes_count,
                                                  "skills_count": skills_count,
                                                  "duration_seconds": (datetime.now() - start_time).total_seconds()
                                              })
                    
                    # 提交事务
                    connection.commit()
                    
                    result = {
                        "success": True,
                        "theater_id": theater_id,
                        "task_id": task_id,
                        "saved_data": {
                            "theater": theater_info.get('title', 'N/A'),
                            "characters": character_count,
                            "script_nodes": nodes_count,
                            "skills": skills_count
                        },
                        "duration_seconds": (datetime.now() - start_time).total_seconds(),
                        "message": f"剧场数据保存成功 (ID: {theater_id})"
                    }
                    
                    logger.info(f"🎉 剧场数据保存成功: Theater ID {theater_id}")
                    return result
                    
                except Exception as e:
                    connection.rollback()
                    error_msg = f"保存数据时发生错误: {str(e)}"
                    logger.error(f"❌ {error_msg}")
                    
                    return {
                        "success": False,
                        "error": error_msg,
                        "task_id": task_id,
                        "duration_seconds": (datetime.now() - start_time).total_seconds()
                    }
                finally:
                    cursor.close()
                    
        except Exception as e:
            error_msg = f"数据库连接失败: {str(e)}"
            logger.error(f"❌ {error_msg}")
            return {
                "success": False,
                "error": error_msg,
                "task_id": task_id,
                "duration_seconds": (datetime.now() - start_time).total_seconds()
            }

    def _save_theater_info(self, cursor, task_id: str, theater_info: Dict) -> int:
        """保存剧场主信息"""
        insert_sql = """
        INSERT INTO theaters (title, description, era, location, plot_summary, features, author, task_id)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
        """
        
        values = (
            theater_info.get('title', '未命名剧场'),
            theater_info.get('description', ''),
            theater_info.get('era', ''),
            theater_info.get('location', ''),
            theater_info.get('plotSummary', ''),
            json.dumps(theater_info.get('features', []), ensure_ascii=False),
            theater_info.get('author', 'AI智能生成器'),
            task_id
        )
        
        cursor.execute(insert_sql, values)
        theater_id = cursor.lastrowid
        logger.info(f"✅ 剧场主信息保存成功 (ID: {theater_id})")
        return theater_id

    def _save_characters(self, cursor, theater_id: int, characters_data: List[Dict]) -> int:
        """保存角色数据"""
        if not characters_data:
            return 0
            
        insert_sql = """
        INSERT INTO characters (theater_id, character_id, name, description, avatar, selected_skills)
        VALUES (%s, %s, %s, %s, %s, %s)
        """
        
        values_list = []
        for char in characters_data:
            values = (
                theater_id,
                char.get('id', ''),
                char.get('name', ''),
                char.get('description', ''),
                char.get('avatar', ''),
                json.dumps(char.get('selectedSkills', []), ensure_ascii=False)
            )
            values_list.append(values)
        
        cursor.executemany(insert_sql, values_list)
        count = cursor.rowcount
        logger.info(f"✅ {count} 个角色保存成功")
        return count

    def _save_script_nodes(self, cursor, theater_id: int, script_nodes: List[Dict]) -> int:
        """保存剧本节点"""
        if not script_nodes:
            return 0
            
        insert_sql = """
        INSERT INTO script_nodes (theater_id, node_id, title, content, character_settings, level, node_order)
        VALUES (%s, %s, %s, %s, %s, %s, %s)
        """
        
        values_list = []
        for i, node in enumerate(script_nodes):
            values = (
                theater_id,
                node.get('id', f'node-{i+1}'),
                node.get('title', ''),
                node.get('content', ''),
                json.dumps(node.get('characterSettings', {}), ensure_ascii=False),
                node.get('level', 0),
                i + 1  # 节点顺序
            )
            values_list.append(values)
        
        cursor.executemany(insert_sql, values_list)
        count = cursor.rowcount
        logger.info(f"✅ {count} 个剧本节点保存成功")
        return count

    def _save_skills_pool(self, cursor, theater_id: int, available_skills: List[str]) -> int:
        """保存技能池"""
        if not available_skills:
            return 0
            
        insert_sql = """
        INSERT INTO skills_pool (theater_id, skill_name, skill_order)
        VALUES (%s, %s, %s)
        ON DUPLICATE KEY UPDATE skill_order = VALUES(skill_order)
        """
        
        values_list = []
        for i, skill in enumerate(available_skills):
            if isinstance(skill, str) and skill.strip():
                values = (theater_id, skill.strip(), i + 1)
                values_list.append(values)
        
        if values_list:
            cursor.executemany(insert_sql, values_list)
            count = len(values_list)
        else:
            count = 0
            
        logger.info(f"✅ {count} 个技能保存成功")
        return count

    def _log_generation_result(self, cursor, task_id: str, theater_id: Optional[int], 
                             stage: str, message: str, metadata: Dict = None):
        """记录生成结果日志"""
        insert_sql = """
        INSERT INTO generation_logs (task_id, theater_id, stage, message, metadata)
        VALUES (%s, %s, %s, %s, %s)
        """
        
        values = (
            task_id,
            theater_id,
            stage,
            message,
            json.dumps(metadata or {}, ensure_ascii=False)
        )
        
        cursor.execute(insert_sql, values)

    async def save_theater_data_async(self, task_id: str, theater_data: Dict[str, Any]) -> Dict[str, Any]:
        """异步保存剧场数据"""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(
            self.executor, 
            self.save_theater_data, 
            task_id, 
            theater_data
        )

    def get_theaters_list(self, limit: int = 20, offset: int = 0) -> Dict:
        """获取剧场列表"""
        try:
            with self.get_connection() as connection:
                cursor = connection.cursor(dictionary=True)
                
                # 获取总数
                cursor.execute("SELECT COUNT(*) as total FROM theaters")
                total = cursor.fetchone()['total']
                
                # 获取剧场列表
                cursor.execute("""
                    SELECT id, title, description, era, location, author, created_at,
                           (SELECT COUNT(*) FROM characters WHERE theater_id = theaters.id) as character_count,
                           (SELECT COUNT(*) FROM script_nodes WHERE theater_id = theaters.id) as node_count,
                           (SELECT COUNT(*) FROM skills_pool WHERE theater_id = theaters.id) as skill_count
                    FROM theaters 
                    ORDER BY created_at DESC 
                    LIMIT %s OFFSET %s
                """, (limit, offset))
                
                theaters = cursor.fetchall()
                
                return {
                    "success": True,
                    "total": total,
                    "theaters": theaters,
                    "limit": limit,
                    "offset": offset
                }
                
        except Exception as e:
            logger.error(f"❌ 获取剧场列表失败: {e}")
            return {
                "success": False,
                "error": str(e),
                "total": 0,
                "theaters": []
            }

    def close(self):
        """关闭资源"""
        if hasattr(self, 'executor'):
            self.executor.shutdown(wait=True)
        logger.info("🔒 MySQL管理器已关闭")

# 创建全局实例
mysql_manager = None

def get_mysql_manager() -> TheaterMySQLManager:
    """获取MySQL管理器实例"""
    global mysql_manager
    if mysql_manager is None:
        mysql_manager = TheaterMySQLManager()
    return mysql_manager

# 测试函数
def test_mysql_connection():
    """测试MySQL连接"""
    try:
        manager = TheaterMySQLManager()
        
        # 测试基本连接
        with manager.get_connection() as connection:
            cursor = connection.cursor()
            cursor.execute("SELECT 1 as test, VERSION() as version, DATABASE() as current_db")
            result = cursor.fetchone()
            cursor.close()
            
        print("✅ MySQL连接测试成功")
        print(f"📊 MySQL版本: {result[1] if result else 'Unknown'}")
        print(f"📊 当前数据库: {result[2] if result else 'Unknown'}")
        
        # 测试获取剧场列表
        theaters_list = manager.get_theaters_list(limit=5)
        print(f"📋 剧场总数: {theaters_list.get('total', 0)}")
        print(f"📋 剧场列表: {len(theaters_list.get('theaters', []))} 条记录")
        
        return True
        
    except Exception as e:
        print(f"❌ MySQL连接测试失败: {e}")
        return False

if __name__ == "__main__":
    print("🧪 运行MySQL模块测试...")
    test_mysql_connection()
